package com.zili.state;

import com.zili.source.ClickSource;
import com.zili.source.Event;
import org.apache.flink.api.common.eventtime.SerializableTimestampAssigner;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.AggregatingState;
import org.apache.flink.api.common.state.AggregatingStateDescriptor;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.MapState;
import org.apache.flink.api.common.state.ReducingState;
import org.apache.flink.api.common.state.ReducingStateDescriptor;
import org.apache.flink.api.common.state.StateTtlConfig;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.Collector;

import java.time.Duration;

/**
 * @author : ranzlupup
 * @since : 2023/3/14 16:54
 */
public class StateTest {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        SingleOutputStreamOperator<Event> stream = env.addSource(new ClickSource())
                .assignTimestampsAndWatermarks(WatermarkStrategy.<Event>forBoundedOutOfOrderness(Duration.ZERO)
                        .withTimestampAssigner(new SerializableTimestampAssigner<Event>() {
                            @Override
                            public long extractTimestamp(Event element, long recordTimestamp) {
                                return element.timestamp;
                            }
                        }));

        stream.keyBy(data -> data.user)
                .flatMap(new MyFlatMap())
                .print();

        env.execute();
    }

    private static class MyFlatMap extends RichFlatMapFunction<Event, String> {
        // 定义状态
        private ValueState<Event> myValueState;
        private ListState<Event> myListState;
        private MapState<String, Long> myMapState;
        private ReducingState<Event> myReduceState;
        private AggregatingState<Event, String> myAggregatingState;
        Long count = 0L;

        @Override
        public void open(Configuration parameters) throws Exception {
            // 创建TTL
            StateTtlConfig stateTtlConfig = StateTtlConfig
                    .newBuilder(Time.seconds(1))
                    .setUpdateType(StateTtlConfig.UpdateType.OnCreateAndWrite)
                    .setStateVisibility(StateTtlConfig.StateVisibility.ReturnExpiredIfNotCleanedUp)
                    .build();

            ValueStateDescriptor<Event> valueStateDescriptor = new ValueStateDescriptor<>("my-state", Event.class);
            myValueState = getRuntimeContext().getState(valueStateDescriptor);

            // 启动TTL
            valueStateDescriptor.enableTimeToLive(stateTtlConfig);

            // myListState = getRuntimeContext().getListState(new ListStateDescriptor<Event>("my-list", Event.class));
            // myMapState = getRuntimeContext().getMapState(new MapStateDescriptor<String, Long>("my-map", String.class, Long.class));
            myReduceState = getRuntimeContext().getReducingState(new ReducingStateDescriptor<Event>("my-reduce", new ReduceFunction<Event>() {
                @Override
                public Event reduce(Event value1, Event value2) throws Exception {
                    return new Event(value1.user, value1.url, value2.timestamp);
                }
            }, Event.class));

            myAggregatingState = getRuntimeContext().getAggregatingState(new AggregatingStateDescriptor<Event, Long, String>("my-agg"
                    , new AggregateFunction<Event, Long, String>() {
                @Override
                public Long createAccumulator() {
                    return 0L;
                }

                @Override
                public Long add(Event value, Long accumulator) {
                    return accumulator + 1;
                }

                @Override
                public String getResult(Long accumulator) {
                    return "count: " + accumulator;
                }

                @Override
                public Long merge(Long a, Long b) {
                    return a + b;
                }
            }
                    , Long.class));
        }

        @Override
        public void flatMap(Event value, Collector<String> out) throws Exception {
            // 访问和更新状态
            System.out.println(myValueState.value());
            // 更新状态
            myValueState.update(value);
            System.out.println("my value" + myValueState.value());

            // myListState.get().forEach(System.out::print);
            // myListState.add(value);
            // System.out.println("my ListValue" + myListState);
            // System.out.println(value.user + "->" + myMapState.get(value.user));
            // myMapState.put(value.user, myMapState.get(value.user) == null ? 1L : myMapState.get(value.user) + 1L);
            // System.out.println(value.user + "->" + myMapState.get(value.user));

            myReduceState.add(value);
            System.out.println("my reducing value: " + myReduceState.get());

            // myAggregatingState.add(value);
            // System.out.println( "my agg value: " + myAggregatingState.get() );
            //
            // count ++;
            // System.out.println("count: " + count);
        }
    }
}
